#pragma once
#include <vector>

#include <algorithm>
#include <array>
#include <random>
#include <cmath>
#include <tuple>

size_t getRandomIdx( size_t max ) {
    static std::random_device r;
    static std::default_random_engine e1( r() );
    std::uniform_int_distribution<size_t> uniform_dist( 0, max );
    return uniform_dist( e1 );
}

size_t getMaxCount( size_t size ) {
    return std::ceil( std::log( size ) / std::log( 3 ) );
}

template <class T = size_t>
struct ScalesN
{
    using DataT = std::vector<T>;
    ScalesN( size_t correctIndex, size_t ballCount )
        : count(getMaxCount(ballCount)), answer( correctIndex ), size(ballCount){
        if( correctIndex > ballCount - 1 ){
            throw std::logic_error{"impossible!"};
        }
        data = std::vector<T>(ballCount, 1);
        data[answer]++;
    };
    ScalesN( const ScalesN & ) = delete;
    ScalesN( ScalesN && ) = delete;

    int getWeight( const std::vector<T> & left, const std::vector<T> & right ) {
        if ( !count )
        { throw std::runtime_error{ "broken" }; }
        count--;
        auto getWeight = [this]( const auto & indexes ) {
            T weight{ 0 };
            for ( const auto & idx : indexes )
            { weight += data[idx]; }
            return weight;
        };
        T leftWeight = getWeight( left ), rightWeight = getWeight( right );
        if ( leftWeight > rightWeight )
        {
            return -1;
        } else if ( rightWeight > leftWeight )
        {
            return 1;
        } else
            return 0;
    }

    size_t getAnswer() const { return answer; }
    size_t getSize() const { return size; }
    size_t getCount() const { return count; }

private:
    size_t count;
    size_t answer;
    const size_t size;
    std::vector<T> data{};
};

namespace eightBalls
{
    // https://www.codewars.com/kata/544047f0cf362503e000036e/train/cpp
    // https://www.codewars.com/kata/54404a06cf36258b08000364/train/cpp
    // https://www.codewars.com/kata/544034f426bc6adda200000e/train/cpp
    using Scales = ScalesN<int>;
    using DataT = Scales::DataT;
    int helper( Scales & scales, const DataT & left, const DataT & right, const DataT & remain ) {
        DataT answer = remain;
        switch ( scales.getWeight( left, right ) )
        {
            case 1:
                answer = right;
                break;
            case -1:
                answer = left;
                break;
        }
        if ( answer.size() == 1 )
            return answer.front();
        return helper( scales, { answer[0] }, { answer[1] }, { answer[2] } );
    }

    unsigned int findBall( Scales scales ) {
        return helper( scales, { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7 } );
    }
} // namespace eightBalls  

namespace dynamicBalls
{
    // https://www.codewars.com/kata/545c4f7682e55d3c6e0011a1
    using Scales = ScalesN<int>;
    using DataT = Scales::DataT;

    auto splitData(const DataT& data){

        DataT left{};
        DataT right{};
        DataT remain{};
        size_t count = (data.size() + 1 )/ 3;
        for( size_t n = 0; n < count; ++n ){
            left.push_back(data[n]);
        }
        for( size_t n = count; n < count * 2; ++n ){
            right.push_back(data[n]);
        }
        for( size_t n = count * 2; n < data.size(); ++n ){
            remain.push_back(data[n]);
        }
        return std::make_tuple(std::move(left), std::move(right), std::move(remain));
    }

    int helper(Scales & scales, const DataT & data){
        if( data.size() == 1 ){
            return data.front();
        }
        auto [left, right, remain] = splitData(data);
        switch( scales.getWeight(left, right) ){
            case 1: return helper(scales, right);
            case -1: return helper(scales, left);
            default: return helper(scales, remain);
        }
    }

    unsigned int findBall( Scales scales ) {
        DataT indexes(scales.getSize());
        std::iota(indexes.begin(), indexes.end(), 0);
        return helper(scales, indexes);
    }

    void test(){
        for ( size_t balls = 3; balls < 512; ++balls ){
            for( size_t idx = 0; idx < balls; ++idx ){
                if( idx != findBall({idx, balls}) ){
                    throw std::runtime_error{"wrong answer!"};
                }
            }
        }
    }
}
